package codingforgreat.class03;

import java.util.*;

public class Class08_DistanceKNodes {
    public static class Node {
        public int value;
        public Node left;
        public Node right;

        public Node(int v) {
            value = v;
        }
    }
    public static List<Node> distanceKNodes(Node root, Node target, int K){
        List<Node> res = new ArrayList<>();
        HashMap<Node,Node> parent = new HashMap<>();
        parent.put(root,null);
        getParentMap(root,parent);
        Queue<Node> queue = new LinkedList<>();
        HashSet<Node> visited = new HashSet<>();
        queue.add(target);
        visited.add(target);
        int nowLevel = 0;
        while (!queue.isEmpty()){
            int size = queue.size();
            while (size-- > 0){
                Node cur = queue.poll();
                if (nowLevel == K) {
                    res.add(cur);
                }
                if (cur.left != null && !visited.contains(cur.left)) {
                    visited.add(cur.left);
                    queue.offer(cur.left);
                }
                if (cur.right != null && !visited.contains(cur.right)) {
                    visited.add(cur.right);
                    queue.offer(cur.right);
                }
                if (parent.get(cur) != null && !visited.contains(parent.get(cur))) {
                    visited.add(parent.get(cur));
                    queue.offer(parent.get(cur));
                }
            }
            nowLevel++;
            if(nowLevel > K){
                break;
            }
        }
        return res;
    }
    public static void getParentMap(Node root,HashMap<Node,Node> parent){
        if(root == null){
            return;
        }
        if(root.left != null){
            parent.put(root.left,root);
            getParentMap(root.left,parent);
        }
        if(root.right != null){
            parent.put(root.right,root);
            getParentMap(root.right,parent);
        }
    }


    public static void main(String[] args) {
        Node n0 = new Node(0);
        Node n1 = new Node(1);
        Node n2 = new Node(2);
        Node n3 = new Node(3);
        Node n4 = new Node(4);
        Node n5 = new Node(5);
        Node n6 = new Node(6);
        Node n7 = new Node(7);
        Node n8 = new Node(8);

        n3.left = n5;
        n3.right = n1;
        n5.left = n6;
        n5.right = n2;
        n1.left = n0;
        n1.right = n8;
        n2.left = n7;
        n2.right = n4;

        Node root = n3;
        Node target = n5;
        int K = 2;

        List<Node> ans = distanceKNodes(root, target, K);
        for (Node o1 : ans) {
            System.out.println(o1.value);
        }

    }

}
